package com.tutu.mybatis.dynamic.datasource.listener;

import com.baomidou.mybatisplus.extension.service.IService;
import com.tutu.mybatis.dynamic.datasource.annotation.DynamicDatasource;
import lombok.NonNull;
import org.springframework.aop.support.AopUtils;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.ApplicationListener;
import org.springframework.context.event.ContextRefreshedEvent;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.*;
import java.util.stream.Collectors;

/**
 * 收集所有标注了<code>@DynamicDatasource</code>的类和方法
 * @see DynamicDatasource
 * @author qyy
 * @date 2021/3/24 14:31
 */
@Component
public class CollectDynamicAnnotationListener implements ApplicationListener<ContextRefreshedEvent>, ApplicationContextAware {

    private ApplicationContext applicationContext;

    public static final Set<String> DYNAMIC_DATASOURCE_SERVICE_CACHE = new HashSet<>();

    @Override
    public void onApplicationEvent(@NonNull ContextRefreshedEvent contextRefreshedEvent) {
        Collection<?> services = applicationContext.getBeansOfType(IService.class).values();
        Set<String> set = services
                .stream()
                .filter(service -> hasDynamicDatasourceAnnotation(service.getClass()))
                .map(service -> {
                    if (AopUtils.isAopProxy(service)) {
                        return AopUtils.getTargetClass(service).getName();
                    }
                    return service.getClass().getName();
                })
                .collect(Collectors.toSet());
        DYNAMIC_DATASOURCE_SERVICE_CACHE.addAll(set);
        List<?> list = services
                .stream()
                .filter(service -> !hasDynamicDatasourceAnnotation(service.getClass()))
                .collect(Collectors.toList());
        for (Object service : list) {
            Class<?> clazz = AopUtils.isAopProxy(service) ? AopUtils.getTargetClass(service) : service.getClass();
            Method[] methods = clazz.getMethods();
            Set<String> methodSet = Arrays.stream(methods)
                    .filter(this::hasDynamicDatasourceAnnotation)
                    .map(method -> clazz.getName() + "#" + method)
                    .collect(Collectors.toSet());
            if (!methodSet.isEmpty()) {
                DYNAMIC_DATASOURCE_SERVICE_CACHE.addAll(methodSet);
            }
        }
    }

    @Override
    public void setApplicationContext(@NonNull ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    private boolean hasDynamicDatasourceAnnotation(Class<?> clazz) {
        return Objects.nonNull(AnnotationUtils.findAnnotation(clazz, DynamicDatasource.class));
    }

    private boolean hasDynamicDatasourceAnnotation(Method method) {
        return Objects.nonNull(AnnotationUtils.findAnnotation(method, DynamicDatasource.class));
    }
}
